{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\kappa(u) = \\left(\\frac{3u(1-u)}{3u^3 + (1-u)^3}\\right)^2  + \\mu\n",
    "$$\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\left\\{\n",
    "\t\\begin{aligned}\n",
    "\t&-\\nabla\\cdot(\\kappa(u)\\nabla u) = f, &x \\in \\Omega, \\\\\n",
    "\t&u = 0, &x \\in \\partial\\Omega\n",
    "\\end{aligned}\n",
    "\\right.\n",
    "\\end{equation}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\tf(x; i^\\star) = 100 \\exp(-50 ||x - x_{c, i^\\star}||_2^2)\\mathbf{1}_{\\Omega_{i^\\star}}(x)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating:: 100%|██████████| 100/100 [03:47<00:00,  2.27s/it]\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from utils import fd_solve_nlinear, force\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "GridSize = 128\n",
    "area = ((0, 0), (1, 1))\n",
    "forces, mus, refs = [], [], []\n",
    "xx, yy = np.meshgrid(\n",
    "    np.linspace(area[0][0], area[1][0], GridSize),\n",
    "    np.linspace(area[0][1], area[1][1], GridSize)\n",
    ")\n",
    "N = 100\n",
    "\n",
    "with tqdm(total=N, desc='Generating:') as pbar:\n",
    "\twhile len(refs) < N:\n",
    "\t\tcenter = np.random.uniform(0.1, 0.9, 2)\n",
    "\t\tmu = np.random.uniform(0.1, 1)\n",
    "\t\tf = force(xx, yy, center)\n",
    "\t\tsol = fd_solve_nlinear(GridSize, area, mu, center)\n",
    "\n",
    "\t\tif not sol is None:\n",
    "\t\t\trefs.append(sol.reshape(GridSize, GridSize))\n",
    "\t\t\tforces.append(f)\n",
    "\t\t\tmus.append(mu * np.ones((GridSize, GridSize)))\n",
    "\t\t\tpbar.update(1)\n",
    "\n",
    "forces = np.stack(forces)\n",
    "mus = np.stack(mus)\n",
    "refs = np.stack(refs)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "p =Path(f'./fit_nlinear/DLData/{GridSize}')\n",
    "if not p.is_dir():\n",
    "\tp.mkdir(parents=True)\n",
    "np.save(f'./fit_nlinear/DLData/{GridSize}/F.npy', forces)\n",
    "np.save(f'./fit_nlinear/DLData/{GridSize}/MU.npy', mus)\n",
    "np.save(f'./fit_nlinear/DLData/{GridSize}/U.npy', refs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from FVM.src.UniformICD import UniformFVM\n",
    "from FVM.src.Problem.test_uniform import *\n",
    "from FVM.src.utils import show_errors\n",
    "\n",
    "# Test zxp case !!!\n",
    "case_ns = list(range(10))\n",
    "print(f\"|{'Method': ^8s}|{'Case ID': ^8s}|{'Mesh Type': ^18s}|{'hmesh': ^11s}|{'error_max': ^11s}|{'ratio': ^7s}|{'error_l2': ^11s}|{'ratio': ^7s}|{'error_h1': ^11s}|{'ratio': ^7s}|\")\n",
    "print(f'|{\":--:\": ^8s}|{\":--:\": ^8s}|{\":--:\": ^18s}|{\":--:\": ^11s}|{\":--:\": ^11s}|{\":--:\": ^7s}|{\":--:\": ^11s}|{\":--:\": ^7s}|{\":--:\": ^11s}|{\":--:\": ^7s}|')\n",
    "for case_n in case_ns:\n",
    "\tmesh_type = 'quadrangle_uniform'\n",
    "\t\t# triangle_classic triangle_random triangle_uniform triangle_kershaw\n",
    "\t\t# quadrangle_uniform quadrangle_random  quadrangle_kershaw\n",
    "\tn_levels = 5\n",
    "\thmeshes = []\n",
    "\terrors_max = []\n",
    "\terrors_l2 = []\n",
    "\terrors_h1 = []\n",
    "\n",
    "\tfor i in range(n_levels):\n",
    "\t\tN = 2**(i+3)\n",
    "\t\tproblem = zxpProblem(N, N, ((0, 0), (1, 1)), case_n, 1e-9)\n",
    "\t\tsolver = UniformFVM(((0, 0), (1, 1)), N, N, problem)\n",
    "\t\tsolver.solve(solver_name=None)\n",
    "\t\thmesh, error_max, error_l2, error_h1 = solver.compute_errors()\n",
    "\n",
    "\t\thmeshes.append(hmesh)\n",
    "\t\terrors_max.append(error_max)\n",
    "\t\terrors_l2.append(error_l2)\n",
    "\t\terrors_h1.append(error_h1)\n",
    "\n",
    "\tfor s in show_errors(hmeshes, errors_max, errors_l2, errors_h1):\n",
    "\t\tprint(f'|{\"ICD\": ^8s}|{case_n: ^8}|{mesh_type: ^18s}|' + s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Problems.NonlinearProblem import *\n",
    "# from plots import my_plot\n",
    "from UniformICD import UniformFVM\n",
    "\n",
    "def fvm_solve_nlinear(GridSize, area, mu, center=(0.5, 0.5), \n",
    "\t\t\t\t\tPicard_maxiter=1000, eps=1e-7):\n",
    "\tu0 = np.random.rand(GridSize, GridSize)\n",
    "\tproblem = NonLinearProblem(None, center, N, area)\n",
    "\tsolver = UniformFVM(area, N, N, problem)\n",
    "\tsolver.solve(solver_name = None)\n",
    "\tu0,  b = solver.ua, solver.b\n",
    "\n",
    "\tfor i in range(Picard_maxiter):\n",
    "\t\tproblem = NonLinearProblem(kappa(u0.reshape(GridSize, GridSize), mu), center, N, area)\n",
    "\n",
    "\t\tsolver = UniformFVM(area, N, N, problem)\n",
    "\t\tsolver.solve(solver_name = None)\n",
    "\t\tnewA = solver.A.tocsr()\n",
    "\t\tnewu = solver.ua\n",
    "\n",
    "\t\t# delta = ((newu - u0)**2 * h**2).sum()\n",
    "\t\tdelta = np.linalg.norm(newu - u0)\n",
    "\t\t# error = ((newA @ u0 - b)**2 * h**2).sum()\n",
    "\t\terror = np.linalg.norm(newA @ u0 - b)\n",
    "\t\tprint(f\"Itr: {i}\\t Delta: {delta:.3e}\\t Error: {error:.3e}\\t\")\n",
    "\n",
    "\t\tif delta <= eps or error <= eps:\n",
    "\t\t\tu = newu.reshape(GridSize, GridSize)\n",
    "\t\t\tbreak\n",
    "\t\telse:\n",
    "\t\t\tu0, A0 = newu, newA\n",
    "\treturn u\n",
    "\n",
    "\n",
    "N = 128\n",
    "mu = 0.1\n",
    "center = (0.1, 0.1)\n",
    "area = ((0, 0), (1, 1))\n",
    "\n",
    "solution = fvm_solve_nlinear(N, area, mu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import segmentation_models_pytorch as smp\n",
    "from SegModel import CNN\n",
    "import torch\n",
    "GridSize = 64\n",
    "net_kwargs={\n",
    "\t\"model_name\": \"varyunet\",\n",
    "\t\"Block\": \"ResBottleNeck\",\n",
    "\t\"planes\": 6,\n",
    "\t\"in_channels\": 1,\n",
    "\t\"classes\": 1,\n",
    "\t\"GridSize\": GridSize,\n",
    "\t\"layer_nums\": [2, 2, 2, 2],\n",
    "\t\"adaptor_nums\": [2, 2, 2, 2],\n",
    "\t\"factor\": 2,\n",
    "\t\"norm_method\": \"layer\",\n",
    "\t\"pool_method\": \"max\",\n",
    "\t\"padding\": \"same\",\n",
    "\t\"padding_mode\": \"zeros\",\n",
    "\t\"end_padding\": \"valid\",\n",
    "\t\"end_padding_mode\": \"zeros\",\n",
    "\t\"act\": \"tanh\"\n",
    "}\n",
    "\n",
    "name = net_kwargs.pop(\"model_name\")\n",
    "net = CNN(name, net_kwargs)\n",
    "\n",
    "x = torch.rand(2, 1, GridSize, GridSize)\n",
    "y = net(x)\n",
    "print(y.shape)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
